Sourced from jax's releases.

JAX v0.7.1

• New features

  • JAX now ships Python 3.14 and 3.14t wheels.
  • JAX now ships Python 3.13t and 3.14t wheels on Mac. Previously we only offered free-threading builds on Linux.

• Changes

  • Exposed jax.set_mesh which acts as a global setter and a context manager. Removed jax.sharding.use_mesh in favor of jax.set_mesh.
  • JAX is now built using CUDA 12.9. All versions of CUDA 12.1 or newer remain supported.
  • jax.lax.dot now implements the general dot product via the optional dimension_numbers argument.

• Deprecations:

  • jax.lax.zeros_like_array is deprecated. Please use jax.numpy.zeros_like instead.
  • Attempting to import jax.experimental.host_callback now results in a DeprecationWarning, and will result in an ImportError starting in JAX v0.8.0. Its APIs have raised NotImplementedError since JAX version 0.4.35.
  • In jax.lax.dot, passing the precision and preferred_element_type arguments by position is deprecated. Pass them by explicit keyword instead.
  • Several dozen internal APIs have been deprecated from jax.interpreters.ad, jax.interpreters.batching, and jax.interpreters.partial_eval; they are used rarely if ever outside JAX itself, and most are deprecated without any public replacement.

Sourced from jax's changelog.

JAX 0.7.1 (August 20, 2025)

• New features

  • JAX now ships Python 3.14 and 3.14t wheels.
  • JAX now ships Python 3.13t and 3.14t wheels on Mac. Previously we only offered free-threading builds on Linux.

• Changes

  • Exposed jax.set_mesh which acts as a global setter and a context manager. Removed jax.sharding.use_mesh in favor of jax.set_mesh.
  • JAX is now built using CUDA 12.9. All versions of CUDA 12.1 or newer remain supported.
  • {func}jax.lax.dot now implements the general dot product via the optional dimension_numbers argument.

• Deprecations:

  • {func}jax.lax.zeros_like_array is deprecated. Please use {func}jax.numpy.zeros_like instead.
  • Attempting to import {mod}jax.experimental.host_callback now results in a DeprecationWarning, and will result in an ImportError starting in JAX v0.8.0. Its APIs have raised NotImplementedError since JAX version 0.4.35.
  • In {func}jax.lax.dot, passing the precision and preferred_element_type arguments by position is deprecated. Pass them by explicit keyword instead.
  • Several dozen internal APIs have been deprecated from {mod}jax.interpreters.ad, {mod}jax.interpreters.batching, and {mod}jax.interpreters.partial_eval; they are used rarely if ever outside JAX itself, and most are deprecated without any public replacement.

JAX 0.7.0 (July 22, 2025)

• New features:

  • Added jax.P which is an alias for jax.sharding.PartitionSpec.
  • Added {func}jax.tree.reduce_associative.
  • The {attr}jax.numpy.ndarray.at indexing methods now support a wrap_negative_indices argument, which defaults to True to match the current behavior ({jax-issue}[#29434](https://github.com/jax-ml/jax/issues/29434)).

• Breaking changes:

  • JAX is migrating from GSPMD to Shardy by default. See the migration guide for more information.
  • JAX autodiff is switching to using direct linearization by default (instead of implementing linearization via JVP and partial eval). See migration guide for more information.
  • jax.stages.OutInfo has been replaced with jax.ShapeDtypeStruct.
  • {func}jax.jit now requires fun to be passed by position, and additional arguments to be passed by keyword. Doing otherwise will result in an error starting in v0.7.x. This raised a DeprecationWarning in v0.6.x.

... (truncated)

• 5712de4 Prepare for JAX release 0.7.1
• 6248205 [Mosaic GPU] Separate hardware-agnostic from hardware-specific passes in BUIL...
• a44869e [Mosaic GPU] Annotate alignment of GMEM pointers to be 256
• 1637155 [Mosaic GPU] Second attempt at fixing test_tcgen05_collective_mma.
• ffc0afb [Mosaic GPU] Attempt no. 2 at fixing the Mac OS build.
• 7ad1508 [Mosaic GPU] Do not predicate collective ops in the lowering.
• cbd0844 [Mosaic GPU] Fold non-gather indices into the gather TMA column index
• 89628e4 Update XLA dependency to use revision http://github.com/openxla/xla/commit/31...
• b09cc05 [Mosaic GPU] Include unistd.h directly in order to fix build on Mac.
• e44e8e6 Merge pull request #31151 from mattjj:fix-error-message
• Additional commits viewable in compare view

You can trigger Dependabot actions by commenting on this PR:

• @dependabot rebase will rebase this PR
• @dependabot recreate will recreate this PR, overwriting any edits that have been made to it
• @dependabot merge will merge this PR after your CI passes on it
• @dependabot squash and merge will squash and merge this PR after your CI passes on it
• @dependabot cancel merge will cancel a previously requested merge and block automerging
• @dependabot reopen will reopen this PR if it is closed
• @dependabot close will close this PR and stop Dependabot recreating it. You can achieve the same result by closing it manually
• @dependabot show <dependency name> ignore conditions will show all of the ignore conditions of the specified dependency
• @dependabot ignore this major version will close this PR and stop Dependabot creating any more for this major version (unless you reopen the PR or upgrade to it yourself)
• @dependabot ignore this minor version will close this PR and stop Dependabot creating any more for this minor version (unless you reopen the PR or upgrade to it yourself)
• @dependabot ignore this dependency will close this PR and stop Dependabot creating any more for this dependency (unless you reopen the PR or upgrade to it yourself)